"""
2024-04-15
增加对核桃派的支持。因为树莓派太贵了。
切换开发版的方法是搜索”树莓派“注释，通过注释找到要切换的语句。


"""


# import RPi.GPIO as GPIO # 树莓派
import board # 核桃派
from  digitalio import DigitalInOut,Direction,Pull # 核桃派

import time
import pyaudio
import wave
import numpy as np

from stt_spark import stt_spark_main # stt_spark_main(audio_file) 可以将音频转换为文字
from tts_spark import tts_spark_main # tts_spark_main(text) 可以将文字转换为音频
from talk_bot import talk_bot # talk_bot(text) 可以返回一个回答


# 录音设置
CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000

# 初始化PyAudio实例
audio = pyaudio.PyAudio()

# 设置GPIO模式和引脚
# GPIO.setmode(GPIO.BCM) # 树莓派
# GPIO.setup(14, GPIO.IN) # 树莓派

key = DigitalInOut(board.PI5) # 核桃派 （注意这里的 【pi5】 是GPIO端口，需要根据接线情况进行适配）
key.direction = Direction.INPUT  #核桃派



# 录音的流变量初始化
stream = None
frames = []

# 开始录音
def start_recording():
    global stream, frames
    stream = audio.open(format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK)
    frames = []

# 结束录音并保存
def stop_recording_and_save(filename):
    global stream
    stream.stop_stream()
    stream.close()
    with wave.open(filename, 'wb') as wf:
        wf.setnchannels(CHANNELS)
        wf.setsampwidth(audio.get_sample_size(FORMAT))
        wf.setframerate(RATE)
        wf.writeframes(b''.join(frames))
    print(f'录音已保存至 {filename}')
    # 在这里调用增益函数
    apply_gain_to_wave(filename, filename, gain=2.0)

# 应用增益到WAV文件
def apply_gain_to_wave(filename, output_filename, gain=20):
    with wave.open(filename, 'rb') as wf:
        params = wf.getparams()
        nchannels, sampwidth, framerate, nframes = params[:4]
        audio_data = wf.readframes(nframes)
    
    if sampwidth == 2:
        data = np.frombuffer(audio_data, dtype=np.int16)
    elif sampwidth == 4:
        data = np.frombuffer(audio_data, dtype=np.int32)
    else:
        raise ValueError("仅支持16位或32位音频数据")

    data = np.int16(data * gain)

    with wave.open(output_filename, 'wb') as wf:
        wf.setnchannels(nchannels)
        wf.setsampwidth(sampwidth)
        wf.setframerate(framerate)
        wf.writeframes(data.tobytes())

# 持续读取音频数据
def continuously_read_data():
    global stream, frames
    data = stream.read(CHUNK, exception_on_overflow=False)
    frames.append(data)

def main_loop():
    key_pressed = False
    try:
        while True:
            # if GPIO.input(14) and not key_pressed: # 树莓派
            if key.value == 1 and not key_pressed: # 核桃派
                key_pressed = True
                print("开始录音...")
                start_recording()
            # elif not GPIO.input(14) and key_pressed: # 树莓派
            elif not key.value == 1 and key_pressed: # 核桃派
                print("停止录音...")
                key_pressed = False
                stop_recording_and_save('./temp/recorded_audio.wav')
                # 移动STT、对话处理和TTS转换逻辑到这里
                try:
                    ask = stt_spark_main('./temp/recorded_audio.wav')
                    print('ask get.')
                    answer = talk_bot(ask)
                    print('anwser get.')
                    tts_spark_main(answer)
                    print('done!')
                except Exception as e:
                    print(f'spark 出错：{e}')
            elif key_pressed:
                continuously_read_data()
            time.sleep(0.01)
    except KeyboardInterrupt:
        print("程序中断。")
    finally:
        # GPIO.cleanup() # 树莓派
        key.deinit() # 核桃派

        audio.terminate()

# 注意：由于环境限制，以下代码行在此环境中不会执行，仅供参考。
if __name__ == "__main__":
    main_loop()
